# 使用预训练数据对BERT模型进行预训练，并保存预训练好的模型和词汇表。
import tokenizers
from transformers import Trainer, TrainingArguments
from transformers import BertTokenizer, LineByLineTextDataset, BertModel, BertConfig, BertForMaskedLM, \
    DataCollatorForLanguageModeling
from bert_config_movie import output_vocab_file

# 设置词汇表文件的路径
vocab_file_dir = './models/vocab.txt'
# 从预训练的词汇表文件加载BERT的分词器（Tokenizer）
tokenizer = BertTokenizer.from_pretrained(vocab_file_dir)

# 将预训练数据转换成适用于预训练的格式
dataset = LineByLineTextDataset(
    tokenizer=tokenizer,  # 指定使用的分词器
    file_path='./working/pretraining_data.txt',  # 指定包含预训练数据的文本文件路径
    block_size=128  # 设置每个样本的最大长度为128个token
)

print('No. of lines: ', len(dataset))

# 定义BERT模型的配置
config = BertConfig(
    vocab_size=90000,  # 词汇表大小
    hidden_size=768,  # 隐层大小
    num_hidden_layers=6,  # BERT模型的隐层层数
    num_attention_heads=12,  # 注意力头数
    max_position_embeddings=512  # 最大位置嵌入
)
# 基于上述配置构建BERT的Masked Language Model（MLM）模型
model = BertForMaskedLM(config)
print('No of parameters: ', model.num_parameters())

# 创建用于语言建模任务的数据收集器
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,  # 指定使用的分词器
    mlm=True,  # 指定为Masked Language Model任务
    mlm_probability=0.15  # 设定mask token的概率
)

# 定义BERT预训练的训练配置
training_args = TrainingArguments(
    output_dir='./working/',  # 设置输出目录路径
    overwrite_output_dir=True,  # 若输出目录已存在，则覆盖
    num_train_epochs=40,  # 训练轮数
    per_device_train_batch_size=8,  # 每个设备的训练批次大小
    save_steps=10_000,  # 每隔10,000步保存一次模型
    save_total_limit=2,  # 最多保存2个模型
)

# 创建用于训练的Trainer对象
trainer = Trainer(
    model=model,  # 指定要训练的BERT模型
    args=training_args,  # 指定训练配置
    data_collator=data_collator,  # 指定数据收集器
    train_dataset=dataset,  # 指定预训练数据集
)

# 执行BERT的预训练
trainer.train()
# 保存模型
trainer.save_model('./models/')
# 将训练好的分词器的词汇表保存到output_vocab_file指定的路径中。
tokenizer.save_vocabulary(output_vocab_file)
